package edu.purdue.cs.ds;

import java.io.*;
import java.net.UnknownHostException;
import java.security.*;
import java.security.spec.*;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

/**
 * Controls the cryptographic operations of Encrypt/Decrypt/Sign/Verify
 * @author jwallrab
 *
 */
public class Cryptography {

	private String local_hostname;
	private PrivateKey privateKey;
	private PublicKey publicKey;
	private SecretKey aesKey;
	private byte[] aesRaw;
	private final int aes_size = 128;
	private Signature dsa;
	
	
	public Cryptography(){
		
		getHostname();
		privateKey = loadPrivateKey();
		publicKey = loadPublicKey(local_hostname);
		initializeSignature();
	}
	
	/**
	 * Return the host name of the local machine
	 */
	public void getHostname(){
		try {
			local_hostname = java.net.InetAddress.getLocalHost().getHostName();
		} catch (UnknownHostException e) {
			e.printStackTrace();
		}
	}
	
	/**
	 * Initialize the DSA Object
	 */
	public void initializeSignature(){
		try{
			dsa = Signature.getInstance("SHA1withRSA");
			dsa.initSign(privateKey);
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}
	}
	
	/**
	 * Load localhost's Private Key from Memory
	 * @param host
	 * @return
	 */
	public PrivateKey loadPrivateKey(){
		
		try{
			
			File f = new File("keys/" + local_hostname + ".pri.der");
			FileInputStream fis = new FileInputStream(f);
			DataInputStream dis = new DataInputStream(fis);
			byte[] keyBytes = new byte[(int)f.length()];
			dis.readFully(keyBytes);
			dis.close();
			
			PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
			KeyFactory kf = KeyFactory.getInstance("RSA");
			return kf.generatePrivate(spec);
			
		}catch(Exception e){
			System.out.println("Error: Could not load Private Key for " + local_hostname);
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Load host's Public Key from Shared Memory
	 * @param host
	 * @return
	 */
	public PublicKey loadPublicKey(String host){
		
		try{
			
			File f = new File("keys/" + host + ".pub.der");
			FileInputStream fis = new FileInputStream(f);
			DataInputStream dis = new DataInputStream(fis);
			byte[] keyBytes = new byte[(int)f.length()];
			dis.readFully(keyBytes);
			dis.close();
			
			X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
			KeyFactory kf = KeyFactory.getInstance("RSA");
			return kf.generatePublic(spec);
			
		}catch(Exception e){
			System.out.println("Error: Could not load Public Key for " + host);
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Encrypt data using the local Public Key
	 * @param data
	 * @return
	 */
	public byte[] encrypt(byte[] data){
		try {
			Cipher cipher = Cipher.getInstance("RSA");
			cipher.init(Cipher.ENCRYPT_MODE, publicKey);
			return cipher.doFinal(data);
		} catch (Exception e){
			System.out.println("Error: Encryption Failed");
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Encrypt data using the host parameter's Public key
	 * @param data
	 * @param host
	 * @return
	 */
	public byte[] encrypt(byte[] data, String host){
		try {
			Cipher cipher = Cipher.getInstance("RSA");
			cipher.init(Cipher.ENCRYPT_MODE, loadPublicKey(host));
			return cipher.doFinal(data);
		} catch (Exception e){
			System.out.println("Error: Encryption Failed");
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Decrypt data using the local Private Key
	 * @param data
	 * @return
	 */
	public byte[] decrypt(byte[] data){
		try {
			Cipher cipher = Cipher.getInstance("RSA");
			cipher.init(Cipher.DECRYPT_MODE, privateKey);
			return cipher.doFinal(data);
		} catch (Exception e){
			System.out.println("Error: Decryption Failed");
			e.printStackTrace();
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Encrypt the data using the AES Key (Must be GENERATED first!)
	 * @param data
	 * @return
	 */
	public byte[] aes_encrypt(byte[] data){
		
		SecretKeySpec skeySpec = new SecretKeySpec(aesRaw,"AES");
		
		try {
			Cipher cipher = Cipher.getInstance("AES");
			cipher.init(Cipher.ENCRYPT_MODE,skeySpec);
			return cipher.doFinal(data);
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Decrypt the data using the AES Key (Must be GENERATED first!)
	 * @param data
	 * @return
	 */
	public byte[] aes_decrypt(byte[] data){
		
		SecretKeySpec skeySpec = new SecretKeySpec(aesRaw,"AES");
		
		try {
			Cipher cipher = Cipher.getInstance("AES");
			cipher.init(Cipher.DECRYPT_MODE,skeySpec);
			return cipher.doFinal(data);
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);
		}
		return null;
	}
	
	/**
	 * Generates a new AES Key
	 * Required before AES Operations can occur, unless the raw
	 * key is read from a message and set using setAESKey(byte[] raw)
	 */
	public void genNewAESKey(){
		try {
			KeyGenerator kgen = KeyGenerator.getInstance("AES");
			kgen.init(aes_size);
			
			aesKey = kgen.generateKey();
			aesRaw = aesKey.getEncoded();
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
			System.exit(1);
		}
		
	}
	
	/**
	 * Sets the raw form of the key for decryption/encryption
	 * @param raw
	 */
	public void setAESKey(byte[] raw){
		aesRaw = raw;
	}
	
	/**
	 * Returns the raw form of the AES Key
	 * @return
	 */
	public byte[] getAESKey(){
		
		if(aesKey == null && aesRaw != null)
			return aesRaw;
		
		return aesKey.getEncoded();
	}
	
	/**
	 * Signs the data
	 * @param data
	 * @return
	 */
	public byte[] sign(byte[] data){
		
		byte[] buffer = new byte[1024];
		byte[] ret = null;
		int len;
		ByteArrayInputStream bais = new ByteArrayInputStream(data);
		BufferedInputStream bufin = new BufferedInputStream(bais);
		try{
			while((len = bufin.read(buffer)) >= 0){
				dsa.update(buffer,0,len);
			}
			bufin.close();
			ret = dsa.sign();
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}
		
		return ret;
	}
	
	/**
	 * Verifies the digital signature of data by host
	 * @param signature
	 * @param data
	 * @param host
	 * @return
	 */
	public boolean verifySignature(byte[] signature, byte[] data, String host){
		
		try{
			Signature sig = Signature.getInstance("SHA1withRSA");
			sig.initVerify(loadPublicKey(host));
			
			ByteArrayInputStream bais = new ByteArrayInputStream(data);
			BufferedInputStream bis = new BufferedInputStream(bais);
			
			byte[] buffer = new byte[1024];
			int len;
			while(bis.available() != 0){
				len = bis.read(buffer);
				sig.update(buffer,0,len);
			}
			bis.close();
			
			return sig.verify(signature);
			
		}catch(Exception e){
			e.printStackTrace();
			System.exit(1);
		}
		return false;
	}
}
